import torch
from torch.utils.data import Dataset
import pandas as pd
import numpy as np
from sklearn.preprocessing import RobustScaler

class EnergyTimeSeriesDataset(Dataset):
    def __init__(self, config, data_path, is_train=True):
        self.config = config
        try:
            self.df = pd.read_excel(data_path)
        except Exception as e:
            print(f"Error reading Excel file: {e}")
            print("Trying with different encoding...")
            try:
                self.df = pd.read_excel(data_path, engine='openpyxl')
            except Exception as e:
                print(f"Still failed: {e}")
                raise
        
        # 分离基本特征和时间特征
        self.base_features = [col for col in self.config.FEATURE_COLUMNS 
                            if col not in ['hour', 'day', 'month', 'weekday']]
        
        # 只检查基本特征列是否存在
        missing_columns = [col for col in self.base_features if col not in self.df.columns]
        if missing_columns:
            print("Warning: Missing columns in dataset:", missing_columns)
            print("Available columns:", self.df.columns.tolist())
            raise ValueError(f"Missing required columns: {missing_columns}")
        
        self.is_train = is_train
        
        # 时间特征工程
        self._add_time_features()
        
        # 数据标准化
        self.scaler = RobustScaler()
        self.data = self._preprocess_data()
        
        # 生成数据片段
        self.segments = self._create_segments()
        
        # 添加数据增强
        if is_train:
            self.data = self.add_noise(self.data)
        
        # 使用RobustScaler而不是StandardScaler
        self.scaler = RobustScaler()
        self.data = self.scaler.fit_transform(self.data)
    
    def _add_time_features(self):
        # 确保timestamp列存在
        if 'timestamp' not in self.df.columns:
            raise ValueError("Missing required 'timestamp' column")
            
        # 转换时间戳并添加时间特征
        self.df['timestamp'] = pd.to_datetime(self.df['timestamp'])
        self.df['hour'] = self.df['timestamp'].dt.hour
        self.df['day'] = self.df['timestamp'].dt.day
        self.df['month'] = self.df['timestamp'].dt.month
        self.df['weekday'] = self.df['timestamp'].dt.weekday
    
    def _preprocess_data(self):
        # 检查无效值
        self.df = self.df.replace([np.inf, -np.inf], np.nan)
        self.df = self.df.fillna(method='ffill').fillna(method='bfill')
        
        # 移除异常值
        for col in self.config.FEATURE_COLUMNS:
            if col in ['hour', 'day', 'month', 'weekday']:
                continue
            Q1 = self.df[col].quantile(0.25)
            Q3 = self.df[col].quantile(0.75)
            IQR = Q3 - Q1
            self.df[col] = self.df[col].clip(lower=Q1 - 1.5*IQR, upper=Q3 + 1.5*IQR)
        
        # 标准化
        data = self.df[self.config.FEATURE_COLUMNS].values
        if self.is_train:
            return self.scaler.fit_transform(data)
        return self.scaler.transform(data)
    
    def _create_segments(self):
        stride = int(self.config.WINDOW_SIZE * (1 - self.config.OVERLAP))
        segments = []
        labels = []
        
        for i in range(0, len(self.data) - self.config.WINDOW_SIZE - self.config.PREDICT_AHEAD + 1, stride):
            segment = self.data[i:i + self.config.WINDOW_SIZE]
            label = self.data[i + self.config.WINDOW_SIZE + self.config.PREDICT_AHEAD - 1]
            segments.append(segment)
            labels.append(label)
            
        return np.array(segments), np.array(labels)
    
    def __len__(self):
        return len(self.segments[0])
    
    def _augment_data(self, segment):
        """数据增强方法"""
        if not self.is_train or np.random.random() > 0.5:
            return segment
        
        # 随机噪声
        noise = np.random.normal(0, 0.01, segment.shape)
        segment = segment + noise
        
        # 随机时间偏移
        shift = np.random.randint(-2, 3)
        if shift > 0:
            segment = np.roll(segment, shift, axis=0)
        
        return segment
    
    def __getitem__(self, idx):
        segment = self.segments[0][idx]
        label = self.segments[1][idx]
        
        # 应用数据增强
        if self.is_train:
            segment = self._augment_data(segment)
        
        return torch.FloatTensor(segment), torch.FloatTensor(label) 
    
    def add_noise(self, data):
        noise = np.random.normal(0, 0.01, data.shape)
        return data + noise 